{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "17da01e4",
   "metadata": {},
   "source": [
    "计算Q时去除基线,基线使用baseline模型估计\n",
    "\n",
    "Q是用蒙特卡洛法估计的,baseline模型的loss就是它的计算结果和Q求mse loss\n",
    "\n",
    "Q在使用时减去baseline模型的计算结果,相当于去基线"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "91162df8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAATrUlEQVR4nO3dfWxTZd8H8G+7rR1jO50brnWuDeSWiAsv6oDtYJ5otDdTFyKyPFGDOg3BMDsizhBdgqiomcE/fMWRO1HwH8TMBI0LinPIiFoYTBfHgEUSTPcgbXl5erpN13Xt9fzhs6OVgXSUXu36/SSHcM51tf2da+t351xnZzUIIQSIiCQwyi6AiDIXA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKSRFkBbtmzBzJkzkZubi8rKSnR1dckqhYgkkRJAH3/8MRobG/HCCy/ghx9+wIIFC1BdXQ2/3y+jHCKSxCDjZtTKykosWrQI7777LgAgGo3Cbrdj7dq1eO6555JdDhFJkp3sFxwdHUV3dzeampr0bUajEU6nE263e8LHhEIhhEIhfT0ajeL8+fMoLi6GwWC46jUTUXyEEBgcHERpaSmMxoufaCU9gM6ePYtIJAKr1Rqz3Wq14vjx4xM+prm5GS+99FIyyiOiBBoYGEBZWdlF25MeQJPR1NSExsZGfV3TNDgcDgwMDEBRFImVEdFEgsEg7HY7CgoKLtkv6QE0Y8YMZGVlwefzxWz3+Xyw2WwTPsZsNsNsNl+wXVEUBhBRCvunKZKkXwUzmUyoqKhAR0eHvi0ajaKjowOqqia7HCKSSMopWGNjI+rq6rBw4UIsXrwYb775JoaHh/H444/LKIeIJJESQA888ADOnDmDjRs3wuv14uabb8aXX355wcQ0EU1tUn4P6EoFg0FYLBZomsY5IKIUdLnvUd4LRkTSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImniDqD9+/dj2bJlKC0thcFgwKeffhrTLoTAxo0bcd1112HatGlwOp34+eefY/qcP38eK1euhKIoKCwsxKpVqzA0NHRFO0JE6SfuABoeHsaCBQuwZcuWCds3b96Mt99+G1u3bsXBgwcxffp0VFdXY2RkRO+zcuVK9PX1ob29HW1tbdi/fz+eeOKJye8FEaUncQUAiF27dunr0WhU2Gw28frrr+vbAoGAMJvN4qOPPhJCCHH06FEBQBw6dEjv88UXXwiDwSBOnTp1Wa+raZoAIDRNu5Lyiegqudz3aELngE6ePAmv1wun06lvs1gsqKyshNvtBgC43W4UFhZi4cKFeh+n0wmj0YiDBw9O+LyhUAjBYDBmIaL0l9AA8nq9AACr1Rqz3Wq16m1erxclJSUx7dnZ2SgqKtL7/F1zczMsFou+2O32RJZNRJKkxVWwpqYmaJqmLwMDA7JLIqIESGgA2Ww2AIDP54vZ7vP59DabzQa/3x/TPjY2hvPnz+t9/s5sNkNRlJiFiNJfQgNo1qxZsNls6Ojo0LcFg0EcPHgQqqoCAFRVRSAQQHd3t95n7969iEajqKysTGQ5RJTisuN9wNDQEE6cOKGvnzx5Ej09PSgqKoLD4cC6devwyiuvYPbs2Zg1axaef/55lJaWYvny5QCAm266CXfffTdWr16NrVu3IhwOo6GhAQ8++CBKS0sTtmNElAbivbz2zTffCAAXLHV1dUKIPy7FP//888JqtQqz2Szuuusu0d/fH/Mc586dEw899JDIz88XiqKIxx9/XAwODib8Eh8RyXG571GDEEJIzL9JCQaDsFgs0DSN80FEKehy36NpcRWMiKYmBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkTdwfy0N0NURGf0fAcwQQ0T82GAyw2Oci25wntzC6qhhAlBJGhwM4uW87RCQMADAYs1C+YgMDaIrjKRilhvT7dChKAAYQpQShf8YlZRIGEKUGHgFlJAYQpYQ0/IBeSgAGEKUIBlAmYgBRauARUEaKK4Cam5uxaNEiFBQUoKSkBMuXL0d/f39Mn5GREbhcLhQXFyM/Px+1tbXw+XwxfTweD2pqapCXl4eSkhKsX78eY2NjV743lLaEiPIgKAPFFUCdnZ1wuVw4cOAA2tvbEQ6HsXTpUgwPD+t9nn76aXz++edobW1FZ2cnfv31V6xYsUJvj0QiqKmpwejoKL7//nt8+OGH2L59OzZu3Ji4vaL0wyOgzCSugN/vFwBEZ2enEEKIQCAgcnJyRGtrq97n2LFjAoBwu91CCCF2794tjEaj8Hq9ep+WlhahKIoIhUKX9bqapgkAQtO0KymfUsjg6RPi0H/qRdfW1aJr62px6D9rxPDZ/5FdFk3S5b5Hr2gOSNM0AEBRUREAoLu7G+FwGE6nU+8zZ84cOBwOuN1uAIDb7ca8efNgtVr1PtXV1QgGg+jr65vwdUKhEILBYMxCU4vg+VdGmnQARaNRrFu3Drfddhvmzp0LAPB6vTCZTCgsLIzpa7Va4fV69T5/DZ/x9vG2iTQ3N8NiseiL3W6fbNmUqgR/ETETTTqAXC4Xjhw5gp07dyayngk1NTVB0zR9GRgYuOqvScklOAeUkSZ1M2pDQwPa2tqwf/9+lJWV6dttNhtGR0cRCARijoJ8Ph9sNpvep6urK+b5xq+Sjff5O7PZDLPZPJlSKV0wgDJSXEdAQgg0NDRg165d2Lt3L2bNmhXTXlFRgZycHHR0dOjb+vv74fF4oKoqAEBVVfT29sLv9+t92tvboSgKysvLr2RfKK0xgDJRXEdALpcLO3bswGeffYaCggJ9zsZisWDatGmwWCxYtWoVGhsbUVRUBEVRsHbtWqiqiqqqKgDA0qVLUV5ejkceeQSbN2+G1+vFhg0b4HK5eJSTwcT43wGijBJXALW0tAAA7rjjjpjt27Ztw2OPPQYAeOONN2A0GlFbW4tQKITq6mq89957et+srCy0tbWhvr4eqqpi+vTpqKurw6ZNm65sTyi98RQsIxlEGs7+BYNBWCwWaJoGRVFkl0MJEPAcwYk9WyCiEQB//kGyvOLrJVdGk3G571HeC0YpIu1+DlICMIAoNaTfgTglAAOIUkIazgRQAjCAKDUIwRDKQAwgSgkCvAyfiRhAlBJE9O8BZAAMUkqhJGIAUUoIBbx/fighgJw8Bdnm6RIromRgAFFKiP7/7/+MMxizYDDy23Oq41eYUpQBPAeb+hhAlJoM+j80hTGAKCUZYICB+TPlMYAoNRl4CpYJGECUupg/Ux4DiFITj4AyAgOIUpKB4ZMRGECUmgwGcBZ66mMAUcpi/Ex9DCBKTZwDyggMIEpRvBk1EzCAKCUZ/vIvTV0MIEpNnIDOCAwgSlEGGBhCUx4DiFITb0bNCAwgSkkGTkJnhLg+GZVossLhMIaHhy/aHhoZiVkfi0SgaRoMxom/Rc1mM6ZNm5bQGin5GECUFN999x0efvjhi7Y/dPts/Pd/zdbXf/rpJyx75hZEohN/UobL5UJTU1PC66TkYgBRUoRCIZw6deqi7drgdRgIzcHp0L9gyT6D4ZE9OHXq1EUDKBgMXq1SKYnimgNqaWnB/PnzoSgKFEWBqqr44osv9PaRkRG4XC4UFxcjPz8ftbW18Pl8Mc/h8XhQU1ODvLw8lJSUYP369RgbG0vM3lDaOh26AceGluB8+Hqc/H0Bjg8t5oelZoC4AqisrAyvvfYauru7cfjwYdx5552477770NfXBwB4+umn8fnnn6O1tRWdnZ349ddfsWLFCv3xkUgENTU1GB0dxffff48PP/wQ27dvx8aNGxO7V5R2hiMWRPUDcgOGIoX8tPgMENcp2LJly2LWX331VbS0tODAgQMoKyvD+++/jx07duDOO+8EAGzbtg033XQTDhw4gKqqKnz11Vc4evQovv76a1itVtx88814+eWX8eyzz+LFF1+EyWRK3J5RWikxeZBrHMRINB9ZhjCuM/3Mz4vPAJOeA4pEImhtbcXw8DBUVUV3dzfC4TCcTqfeZ86cOXA4HHC73aiqqoLb7ca8efNgtVr1PtXV1aivr0dfXx9uueWWuGo4fvw48vPzJ7sLlEQej+eS7f0//4hweBPOjV6P/Oz/xdjQz5c8Ajp79iyOHj2a2CIpYYaGhi6rX9wB1NvbC1VVMTIygvz8fOzatQvl5eXo6emByWRCYWFhTH+r1Qqv1wsA8Hq9MeEz3j7edjGhUAihUEhfH5+A1DSN80dp4lKX4AGg54QXPScu/j3wd6FQCIFA4Aqroqvln77e4+IOoBtvvBE9PT3QNA2ffPIJ6urq0NnZGXeB8WhubsZLL710wfbKykooinJVX5sSY3BwMKHPd/3112PJkiUJfU5KnMu9Shn3b0KbTCbccMMNqKioQHNzMxYsWIC33noLNpsNo6OjF/xU8vl8sNlsAACbzXbBVbHx9fE+E2lqaoKmafoyMDAQb9lElIKu+FaMaDSKUCiEiooK5OTkoKOjQ2/r7++Hx+OBqqoAAFVV0dvbC7/fr/dpb2+HoigoLy+/6GuYzWb90v/4QkTpL65TsKamJtxzzz1wOBwYHBzEjh07sG/fPuzZswcWiwWrVq1CY2MjioqKoCgK1q5dC1VVUVVVBQBYunQpysvL8cgjj2Dz5s3wer3YsGEDXC4XzGbzVdlBIkpdcQWQ3+/Ho48+itOnT8NisWD+/PnYs2cP/v3vfwMA3njjDRiNRtTW1iIUCqG6uhrvvfee/visrCy0tbWhvr4eqqpi+vTpqKurw6ZNmxK7V5RysrOzE3rkyh9YU4NBiPT7ZYtgMAiLxQJN03g6liZGRkZw5syZhD1fQUHBBVdcKXVc7nuU94JRUuTm5sJut8sug1IM/x4QEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISBoGEBFJwwAiImkYQEQkDQOIiKRhABGRNAwgIpKGAURE0jCAiEgaBhARScMAIiJpGEBEJA0DiIikYQARkTQMICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikiZbdgGTIYQAAASDQcmVENFExt+b4+/Vi0nLADp37hwAwG63S66EiC5lcHAQFovlou1pGUBFRUUAAI/Hc8mdo1jBYBB2ux0DAwNQFEV2OWmBYzY5QggMDg6itLT0kv3SMoCMxj+mriwWC78pJkFRFI5bnDhm8bucgwNOQhORNAwgIpImLQPIbDbjhRdegNlsll1KWuG4xY9jdnUZxD9dJyMiukrS8giIiKYGBhARScMAIiJpGEBEJE1aBtCWLVswc+ZM5ObmorKyEl1dXbJLkqa5uRmLFi1CQUEBSkpKsHz5cvT398f0GRkZgcvlQnFxMfLz81FbWwufzxfTx+PxoKamBnl5eSgpKcH69esxNjaWzF2R5rXXXoPBYMC6dev0bRyzJBFpZufOncJkMokPPvhA9PX1idWrV4vCwkLh8/lklyZFdXW12LZtmzhy5Ijo6ekR9957r3A4HGJoaEjvs2bNGmG320VHR4c4fPiwqKqqEkuWLNHbx8bGxNy5c4XT6RQ//vij2L17t5gxY4ZoamqSsUtJ1dXVJWbOnCnmz58vnnrqKX07xyw50i6AFi9eLFwul74eiUREaWmpaG5ullhV6vD7/QKA6OzsFEIIEQgERE5OjmhtbdX7HDt2TAAQbrdbCCHE7t27hdFoFF6vV+/T0tIiFEURoVAouTuQRIODg2L27Nmivb1d3H777XoAccySJ61OwUZHR9Hd3Q2n06lvMxqNcDqdcLvdEitLHZqmAfjzht3u7m6Ew+GYMZszZw4cDoc+Zm63G/PmzYPVatX7VFdXIxgMoq+vL4nVJ5fL5UJNTU3M2AAcs2RKq5tRz549i0gkEvNFBwCr1Yrjx49Lqip1RKNRrFu3Drfddhvmzp0LAPB6vTCZTCgsLIzpa7Va4fV69T4Tjel421S0c+dO/PDDDzh06NAFbRyz5EmrAKJLc7lcOHLkCL799lvZpaS0gYEBPPXUU2hvb0dubq7scjJaWp2CzZgxA1lZWRdcjfD5fLDZbJKqSg0NDQ1oa2vDN998g7KyMn27zWbD6OgoAoFATP+/jpnNZptwTMfbppru7m74/X7ceuutyM7ORnZ2Njo7O/H2228jOzsbVquVY5YkaRVAJpMJFRUV6Ojo0LdFo1F0dHRAVVWJlckjhEBDQwN27dqFvXv3YtasWTHtFRUVyMnJiRmz/v5+eDwefcxUVUVvby/8fr/ep729HYqioLy8PDk7kkR33XUXent70dPToy8LFy7EypUr9f9zzJJE9ix4vHbu3CnMZrPYvn27OHr0qHjiiSdEYWFhzNWITFJfXy8sFovYt2+fOH36tL789ttvep81a9YIh8Mh9u7dKw4fPixUVRWqqurt45eUly5dKnp6esSXX34prr322oy6pPzXq2BCcMySJe0CSAgh3nnnHeFwOITJZBKLFy8WBw4ckF2SNAAmXLZt26b3+f3338WTTz4prrnmGpGXlyfuv/9+cfr06Zjn+eWXX8Q999wjpk2bJmbMmCGeeeYZEQ6Hk7w38vw9gDhmycE/x0FE0qTVHBARTS0MICKShgFERNIwgIhIGgYQEUnDACIiaRhARCQNA4iIpGEAEZE0DCAikoYBRETSMICISJr/AxX3fSf8CH1lAAAAAElFTkSuQmCC",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import gym\n",
    "\n",
    "\n",
    "#定义环境\n",
    "class MyWrapper(gym.Wrapper):\n",
    "\n",
    "    def __init__(self):\n",
    "        env = gym.make('CartPole-v1', render_mode='rgb_array')\n",
    "        super().__init__(env)\n",
    "        self.env = env\n",
    "        self.step_n = 0\n",
    "\n",
    "    def reset(self):\n",
    "        state, _ = self.env.reset()\n",
    "        self.step_n = 0\n",
    "        return state\n",
    "\n",
    "    def step(self, action):\n",
    "        state, reward, terminated, truncated, info = self.env.step(action)\n",
    "        over = terminated or truncated\n",
    "\n",
    "        #限制最大步数\n",
    "        self.step_n += 1\n",
    "        if self.step_n >= 200:\n",
    "            over = True\n",
    "        \n",
    "        #没坚持到最后,扣分\n",
    "        if over and self.step_n < 200:\n",
    "            reward = -1000\n",
    "\n",
    "        return state, reward, over\n",
    "\n",
    "    #打印游戏图像\n",
    "    def show(self):\n",
    "        from matplotlib import pyplot as plt\n",
    "        plt.figure(figsize=(3, 3))\n",
    "        plt.imshow(self.env.render())\n",
    "        plt.show()\n",
    "\n",
    "\n",
    "env = MyWrapper()\n",
    "\n",
    "env.reset()\n",
    "\n",
    "env.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ecfbe912",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=2, bias=True)\n",
       "   (5): Softmax(dim=1)\n",
       " ),\n",
       " Sequential(\n",
       "   (0): Linear(in_features=4, out_features=64, bias=True)\n",
       "   (1): ReLU()\n",
       "   (2): Linear(in_features=64, out_features=64, bias=True)\n",
       "   (3): ReLU()\n",
       "   (4): Linear(in_features=64, out_features=1, bias=True)\n",
       " ))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "\n",
    "#定义模型,计算每个动作的概率\n",
    "model_action = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 2),\n",
    "    torch.nn.Softmax(dim=1),\n",
    ")\n",
    "\n",
    "#基线模型,评估state的价值\n",
    "model_baseline = torch.nn.Sequential(\n",
    "    torch.nn.Linear(4, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 64),\n",
    "    torch.nn.ReLU(),\n",
    "    torch.nn.Linear(64, 1),\n",
    ")\n",
    "\n",
    "model_action, model_baseline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "84cbf0ff",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "d:\\appDir\\python3.10\\lib\\site-packages\\gym\\utils\\passive_env_checker.py:233: DeprecationWarning: `np.bool8` is a deprecated alias for `np.bool_`.  (Deprecated NumPy 1.24)\n",
      "  if not isinstance(terminated, (bool, np.bool8)):\n",
      "C:\\Users\\Administrator\\AppData\\Local\\Temp\\ipykernel_412\\3856618961.py:30: UserWarning: Creating a tensor from a list of numpy.ndarrays is extremely slow. Please consider converting the list to a single numpy.ndarray with numpy.array() before converting to a tensor. (Triggered internally at ..\\torch\\csrc\\utils\\tensor_new.cpp:248.)\n",
      "  state = torch.FloatTensor(state).reshape(-1, 4)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "-987.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from IPython import display\n",
    "import random\n",
    "\n",
    "\n",
    "#玩一局游戏并记录数据\n",
    "def play(show=False):\n",
    "    state = []\n",
    "    action = []\n",
    "    reward = []\n",
    "\n",
    "    s = env.reset()\n",
    "    o = False\n",
    "    while not o:\n",
    "        #根据概率采样\n",
    "        prob = model_action(torch.FloatTensor(s).reshape(1, 4))[0].tolist()\n",
    "        a = random.choices(range(2), weights=prob, k=1)[0]\n",
    "\n",
    "        ns, r, o = env.step(a)\n",
    "\n",
    "        state.append(s)\n",
    "        action.append(a)\n",
    "        reward.append(r)\n",
    "\n",
    "        s = ns\n",
    "\n",
    "        if show:\n",
    "            display.clear_output(wait=True)\n",
    "            env.show()\n",
    "\n",
    "    state = torch.FloatTensor(state).reshape(-1, 4)\n",
    "    action = torch.LongTensor(action).reshape(-1, 1)\n",
    "    reward = torch.FloatTensor(reward).reshape(-1, 1)\n",
    "\n",
    "    return state, action, reward, reward.sum().item()\n",
    "\n",
    "\n",
    "state, action, reward, reward_sum = play()\n",
    "\n",
    "reward_sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "14ac0e81",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer_action = torch.optim.Adam(model_action.parameters(), lr=5e-3)\n",
    "optimizer_baseline = torch.optim.Adam(model_baseline.parameters(), lr=5e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "0342b2ef",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([14, 1])"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def get_value(reward):\n",
    "    #计算当前state的价值,其实就是Q(state,action),这里是用蒙特卡洛法估计的\n",
    "    value = []\n",
    "    for i in range(len(reward)):\n",
    "        s = 0\n",
    "        for j in range(i, len(reward)):\n",
    "            s += reward[j] * 0.99**(j - i)\n",
    "        value.append(s)\n",
    "\n",
    "    return torch.FloatTensor(value).reshape(-1, 1)\n",
    "\n",
    "\n",
    "value = get_value(reward)\n",
    "\n",
    "value.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "99cc52b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([14, 1])"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#训练baseline模型\n",
    "def train_baseline(state, value):\n",
    "    baseline = model_baseline(state)\n",
    "\n",
    "    loss = torch.nn.functional.mse_loss(baseline, value)\n",
    "    loss.backward()\n",
    "    optimizer_baseline.step()\n",
    "    optimizer_baseline.zero_grad()\n",
    "\n",
    "    return baseline.detach()\n",
    "\n",
    "\n",
    "baseline = train_baseline(state, value)\n",
    "\n",
    "baseline.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "42024b65",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-635.6792602539062"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#训练action模型\n",
    "def train_action(state, action, value, baseline):\n",
    "    #重新计算动作的概率\n",
    "    prob = model_action(state).gather(dim=1, index=action)\n",
    "\n",
    "    #求Q最大的导函数 -> partial value / partial action\n",
    "    #注意这里的Q使用前要去基线,这也是baseline模型存在的意义\n",
    "    prob = (prob + 1e-8).log() * (value - baseline)\n",
    "    for i in range(len(prob)):\n",
    "        prob[i] = prob[i] * 0.99**i\n",
    "    loss = -prob.mean()\n",
    "\n",
    "    loss.backward()\n",
    "    optimizer_action.step()\n",
    "    optimizer_action.zero_grad()\n",
    "\n",
    "    return loss.item()\n",
    "\n",
    "\n",
    "train_action(state, action, value, baseline)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "dbe2cd7a",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 -485.3520202636719 -974.15\n"
     ]
    }
   ],
   "source": [
    "#训练\n",
    "def train():\n",
    "    model_action.train()\n",
    "    model_baseline.train()\n",
    "\n",
    "    #训练N局\n",
    "    for epoch in range(100):\n",
    "\n",
    "        #一个epoch最少玩N步\n",
    "        steps = 0\n",
    "        while steps < 200:\n",
    "\n",
    "            #玩一局游戏,得到数据\n",
    "            state, action, reward, _ = play()\n",
    "            steps += len(state)\n",
    "\n",
    "            #训练两个模型\n",
    "            value = get_value(reward)\n",
    "            baseline = train_baseline(state, value)\n",
    "            loss = train_action(state, action, value, baseline)\n",
    "\n",
    "        if epoch % 100 == 0:\n",
    "            test_result = sum([play()[-1] for _ in range(20)]) / 20\n",
    "            print(epoch, loss, test_result)\n",
    "\n",
    "\n",
    "train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a101d0e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAASAAAADMCAYAAADTcn7NAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAA9hAAAPYQGoP6dpAAAT4klEQVR4nO3df2xT97kG8MdOYhOSHIckjd3cxAJpqDTix7YA4Qzttlo9si6qyprqbhPqsopLVeag0kxIy1VLL72VUtGr264bhD+mQaWrjiqV6NSUgiJTwioMKWmjhQDZprVLbsF2IctxkhInjt/7R5ezHgg0ToK/GD8f6Uic831tv+cQPzk/fGKbiAiIiBSwq26AiDIXA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRRFkB79uzB4sWLsWDBAlRXV6Ozs1NVK0SkiJIAeuONN9DY2IjnnnsOH374IVatWoWamhpEIhEV7RCRIjYVN6NWV1djzZo1+PWvfw0ASCQSqKiowLZt2/CLX/wi1e0QkSLZqX7B8fFxdHV1oampyVxmt9vh8/kQDAanfUwsFkMsFjPnE4kEBgcHUVxcDJvNdst7JqLkiAiGh4dRVlYGu/3GB1opD6DLly9jcnISbrfbstztduPChQvTPqa5uRm7du1KRXtENI8GBgZQXl5+w/GUB9BsNDU1obGx0Zw3DANerxcDAwPQNE1hZ0Q0nWg0ioqKChQUFNy0LuUBVFJSgqysLITDYcvycDgMj8cz7WOcTiecTud1yzVNYwAR3ca+6hRJyq+CORwOVFVVIRAImMsSiQQCgQB0XU91O0SkkJJDsMbGRtTX12P16tVYu3YtXnnlFYyOjuLxxx9X0Q4RKaIkgH74wx/is88+w86dOxEKhfD1r38dR44cue7ENBHd2ZR8DmiuotEoXC4XDMPgOSCi29BM36O8F4yIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKJB1AJ06cwEMPPYSysjLYbDa89dZblnERwc6dO3H33XcjNzcXPp8Pf/7zny01g4OD2LRpEzRNQ2FhITZv3oyRkZE5rQgRpZ+kA2h0dBSrVq3Cnj17ph3fvXs3Xn31Vezbtw+nT59GXl4eampqMDY2ZtZs2rQJvb29aG9vR1tbG06cOIEnnnhi9mtBROlJ5gCAHDp0yJxPJBLi8XjkpZdeMpcNDQ2J0+mU3/3udyIicu7cOQEgH3zwgVnz7rvvis1mk08//XRGr2sYhgAQwzDm0j4R3SIzfY/O6zmgjz/+GKFQCD6fz1zmcrlQXV2NYDAIAAgGgygsLMTq1avNGp/PB7vdjtOnT0/7vLFYDNFo1DIRUfqb1wAKhUIAALfbbVnudrvNsVAohNLSUst4dnY2ioqKzJprNTc3w+VymVNFRcV8tk1EiqTFVbCmpiYYhmFOAwMDqlsionkwrwHk8XgAAOFw2LI8HA6bYx6PB5FIxDIej8cxODho1lzL6XRC0zTLRETpb14DaMmSJfB4PAgEAuayaDSK06dPQ9d1AICu6xgaGkJXV5dZc+zYMSQSCVRXV89nO0R0m8tO9gEjIyP4y1/+Ys5//PHH6O7uRlFREbxeL7Zv344XXngBS5cuxZIlS/Dss8+irKwMGzduBADce++9+N73voctW7Zg3759mJiYQENDA370ox+hrKxs3laMiNJAspfX3nvvPQFw3VRfXy8iX1yKf/bZZ8XtdovT6ZQHHnhA+vr6LM9x5coV+fGPfyz5+fmiaZo8/vjjMjw8PO+X+IhIjZm+R20iIgrzb1ai0ShcLhcMw+D5IKLb0Ezfo2lxFYyI7kwMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEiZpL+Wh+4sw5f+hPGRv5vzTlcp8u5aDJvNprAryhQMoAwX+mMAQ598ZM7fde+3kXfXYnUNUUbhIRhZSCKhugXKIAwgshBJ4IvvmiS69RhAZCGSYP5QyjCAyCrBPSBKHQYQWXxxCEaUGgwgshBJQIR7QJQaSQVQc3Mz1qxZg4KCApSWlmLjxo3o6+uz1IyNjcHv96O4uBj5+fmoq6tDOBy21PT396O2thYLFy5EaWkpduzYgXg8Pve1obnjVTBKoaQCqKOjA36/H6dOnUJ7ezsmJiawYcMGjI6OmjVPP/003n77bbS2tqKjowMXL17EI488Yo5PTk6itrYW4+PjOHnyJF577TUcOHAAO3funL+1olnjIRillMxBJBIRANLR0SEiIkNDQ5KTkyOtra1mzfnz5wWABINBERE5fPiw2O12CYVCZk1LS4tomiaxWGxGr2sYhgAQwzDm0j6JyJ+O7JXOfVvM6cI7r8hkfFx1W5TmZvoendM5IMMwAABFRUUAgK6uLkxMTMDn85k1y5Ytg9frRTAYBAAEg0GsWLECbrfbrKmpqUE0GkVvb++0rxOLxRCNRi0T3SIJXoan1Jl1ACUSCWzfvh3r16/H8uXLAQChUAgOhwOFhYWWWrfbjVAoZNZ8OXymxqfGptPc3AyXy2VOFRUVs22brmGzW+/5SiQmIUwgSpFZB5Df78fZs2dx8ODB+exnWk1NTTAMw5wGBgZu+WtmitxF/2KZj0U/g8QnFHVDmWZWN6M2NDSgra0NJ06cQHl5ubnc4/FgfHwcQ0NDlr2gcDgMj8dj1nR2dlqeb+oq2VTNtZxOJ5xO52xapa9gz86xzEtiEjwGo1RJag9IRNDQ0IBDhw7h2LFjWLJkiWW8qqoKOTk5CAQC5rK+vj709/dD13UAgK7r6OnpQSQSMWva29uhaRoqKyvnsi40CzZbluoWKIMltQfk9/vx+uuv4/e//z0KCgrMczYulwu5ublwuVzYvHkzGhsbUVRUBE3TsG3bNui6jnXr1gEANmzYgMrKSjz22GPYvXs3QqEQnnnmGfj9fu7lKGCz87OopE5SAdTS0gIAuP/++y3L9+/fj5/+9KcAgJdffhl2ux11dXWIxWKoqanB3r17zdqsrCy0tbVh69at0HUdeXl5qK+vx/PPPz+3NaHZsXMPiNSxiaTf5+6j0ShcLhcMw4CmaarbSWuRcyfwtz/8rzmfnathxb/9J7IX5CvsitLdTN+j3P/OcDwEI5X405fhbDwEI4UYQBnOZuOPAKnDn75Mxz0gUogBlOF4DohU4k9fhuMHEUklBlCGm+4LCNPwkxmUphhAdA0BGECUIgwgshL+VURKHQYQXUMYQJQyDCCyEgEYQJQiDCCyEPD74Sl1GEB0De4BUeowgMhKeBmeUocBRNcQHoJRyjCAyIonoSmFGEBkIeDngCh1GECZ7tpbMUT+8c0YRLceAyjDOfIWIcuRa84n4jHEhi8r7IgyCQMow9mysoFr/igZ94AoVRhAGc5ms017RzxRKjCAMp3NDoABRGowgDKczWa//kQ0UYowgDIdD8FIoaS+GZXS0+joKCYmJqYdi18dRiJhvfXi888/x9DQ0A2fr6CgAFlZ/FOuNHcMoAzQ2NiId955Z9oxV54D//3v61Fc8M9L8U1N/4GOnk+nrXc4HDh69CiWLl16S3qlzMIAygCDg4P49NPpA2U034nRcScujXwbVxMF8C44h7///fgN651O5w33poiSldQ5oJaWFqxcuRKapkHTNOi6jnfffdccHxsbg9/vR3FxMfLz81FXV4dwOGx5jv7+ftTW1mLhwoUoLS3Fjh07EI/H52dtKGnxSTv+OPyv+L/YMlyZKEfPyP0YnLhbdVuUIZIKoPLycrz44ovo6urCmTNn8J3vfAcPP/wwent7AQBPP/003n77bbS2tqKjowMXL17EI488Yj5+cnIStbW1GB8fx8mTJ/Haa6/hwIED2Llz5/yuFc1YXGwYibswdSk+Lg6MJfLVNkWZQ+Zo0aJF8pvf/EaGhoYkJydHWltbzbHz588LAAkGgyIicvjwYbHb7RIKhcyalpYW0TRNYrHYjF/TMAwBIIZhzLX9jPDoo48K/nGf6bXTAqdDdje9JP/1winZ9UKn/E/zEblvddUN651Op/T29qpeJbrNzfQ9OutzQJOTk2htbcXo6Ch0XUdXVxcmJibg8/nMmmXLlsHr9SIYDGLdunUIBoNYsWIF3G63WVNTU4OtW7eit7cX3/jGN5Lq4cKFC8jP52/rrxKNRm84Fo/HcfL9AxjLfh9jiTyUOv6GgYt/vWG9iOCvf73xOBEAjIyMzKgu6QDq6emBrusYGxtDfn4+Dh06hMrKSnR3d8PhcKCwsNBS73a7EQqFAAChUMgSPlPjU2M3EovFEIvFzPmpN5RhGDx/NAM3O2kcn0zgrT/0Auid0XOJCIaHh296mZ5odHR0RnVJB9A999yD7u5uGIaBN998E/X19ejo6Ei6wWQ0Nzdj165d1y2vrq6Gpmm39LXvBMXFxfP2XHa7HatWrUJlZeW8PSfdeW621/1lSX8S2uFw4Gtf+xqqqqrQ3NyMVatW4Ze//CU8Hg/Gx8ev+80YDofh8XgAAB6P57qrYlPzUzXTaWpqgmEY5jQwMJBs20R0G5rzrRiJRAKxWAxVVVXIyclBIBAwx/r6+tDf3w9d1wEAuq6jp6cHkUjErGlvb4emaTf9jep0Os1L/1MTEaW/pA7Bmpqa8OCDD8Lr9WJ4eBivv/46jh8/jqNHj8LlcmHz5s1obGxEUVERNE3Dtm3boOs61q1bBwDYsGEDKisr8dhjj2H37t0IhUJ45pln4Pf74XQ6b8kKEtHtK6kAikQi+MlPfoJLly7B5XJh5cqVOHr0KL773e8CAF5++WXY7XbU1dUhFouhpqYGe/fuNR+flZWFtrY2bN26FbquIy8vD/X19Xj++efnd63IIjc3d972Gp1OJ+8Do3ljE0m/L4GKRqNwuVwwDIOHYzNw+fJlXL16dV6ey2azwe12IycnZ16ej+5MM32P8l6wDFBSUqK6BaJp8e8BEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyjCAiEgZBhARKcMAIiJlGEBEpAwDiIiUYQARkTIMICJShgFERMowgIhIGQYQESnDACIiZRhARKQMA4iIlGEAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImWyVTcwGyICAIhGo4o7IaLpTL03p96rN5KWAXTlyhUAQEVFheJOiOhmhoeH4XK5bjielgFUVFQEAOjv77/pypFVNBpFRUUFBgYGoGma6nbSArfZ7IgIhoeHUVZWdtO6tAwgu/2LU1cul4s/FLOgaRq3W5K4zZI3k50DnoQmImUYQESkTFoGkNPpxHPPPQen06m6lbTC7ZY8brNbyyZfdZ2MiOgWScs9ICK6MzCAiEgZBhARKcMAIiJl0jKA9uzZg8WLF2PBggWorq5GZ2en6paUaW5uxpo1a1BQUIDS0lJs3LgRfX19lpqxsTH4/X4UFxcjPz8fdXV1CIfDlpr+/n7U1tZi4cKFKC0txY4dOxCPx1O5Ksq8+OKLsNls2L59u7mM2yxFJM0cPHhQHA6H/Pa3v5Xe3l7ZsmWLFBYWSjgcVt2aEjU1NbJ//345e/asdHd3y/e//33xer0yMjJi1jz55JNSUVEhgUBAzpw5I+vWrZNvfetb5ng8Hpfly5eLz+eTjz76SA4fPiwlJSXS1NSkYpVSqrOzUxYvXiwrV66Up556ylzObZYaaRdAa9euFb/fb85PTk5KWVmZNDc3K+zq9hGJRASAdHR0iIjI0NCQ5OTkSGtrq1lz/vx5ASDBYFBERA4fPix2u11CoZBZ09LSIpqmSSwWS+0KpNDw8LAsXbpU2tvb5b777jMDiNssddLqEGx8fBxdXV3w+XzmMrvdDp/Ph2AwqLCz24dhGAD+ecNuV1cXJiYmLNts2bJl8Hq95jYLBoNYsWIF3G63WVNTU4NoNIre3t4Udp9afr8ftbW1lm0DcJulUlrdjHr58mVMTk5a/tMBwO1248KFC4q6un0kEgls374d69evx/LlywEAoVAIDocDhYWFllq3241QKGTWTLdNp8buRAcPHsSHH36IDz744LoxbrPUSasAopvz+/04e/Ys3n//fdWt3NYGBgbw1FNPob29HQsWLFDdTkZLq0OwkpISZGVlXXc1IhwOw+PxKOrq9tDQ0IC2tja89957KC8vN5d7PB6Mj49jaGjIUv/lbebxeKbdplNjd5quri5EIhF885vfRHZ2NrKzs9HR0YFXX30V2dnZcLvd3GYpklYB5HA4UFVVhUAgYC5LJBIIBALQdV1hZ+qICBoaGnDo0CEcO3YMS5YssYxXVVUhJyfHss36+vrQ399vbjNd19HT04NIJGLWtLe3Q9M0VFZWpmZFUuiBBx5AT08Puru7zWn16tXYtGmT+W9usxRRfRY8WQcPHhSn0ykHDhyQc+fOyRNPPCGFhYWWqxGZZOvWreJyueT48eNy6dIlc/r888/NmieffFK8Xq8cO3ZMzpw5I7qui67r5vjUJeUNGzZId3e3HDlyRO66666MuqT85atgItxmqZJ2ASQi8qtf/Uq8Xq84HA5Zu3atnDp1SnVLygCYdtq/f79Zc/XqVfnZz34mixYtkoULF8oPfvADuXTpkuV5PvnkE3nwwQclNzdXSkpK5Oc//7lMTEykeG3UuTaAuM1Sg3+Og4iUSatzQER0Z2EAEZEyDCAiUoYBRETKMICISBkGEBEpwwAiImUYQESkDAOIiJRhABGRMgwgIlKGAUREyvw/YEWp71vS5WUAAAAASUVORK5CYII=",
      "text/plain": [
       "<Figure size 300x300 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "200.0"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "play(True)[-1]"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
